package com.tianji.aigc.service.impl;

import cn.hutool.core.date.DateUtil;
import cn.hutool.core.map.MapUtil;
import cn.hutool.core.util.IdUtil;
import com.tianji.aigc.config.SystemPromptConfig;
import com.tianji.aigc.config.ToolResultHolder;
import com.tianji.aigc.constants.Constant;
import com.tianji.aigc.enums.ChatEventTypeEnum;
import com.tianji.aigc.service.ChatService;
import com.tianji.aigc.service.ChatSessionService;
import com.tianji.aigc.vo.ChatEventVO;
import com.tianji.common.utils.UserContext;
import lombok.RequiredArgsConstructor;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.client.advisor.AbstractChatMemoryAdvisor;
import org.springframework.ai.chat.client.advisor.QuestionAnswerAdvisor;
import org.springframework.ai.chat.memory.ChatMemory;
import org.springframework.ai.chat.messages.AssistantMessage;
import org.springframework.ai.vectorstore.SearchRequest;
import org.springframework.ai.vectorstore.VectorStore;
import reactor.core.publisher.Flux;

import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

/**
 * @author wys
 * @version V1.0
 * @date 2025-04-21 16:20
 */
//@Service
@RequiredArgsConstructor
public class ChatServiceImpl implements ChatService {

    private final ChatClient dashScopeChatClient;

    private final SystemPromptConfig systemPromptConfig;

    private final ChatMemory chatMemory;

    private final VectorStore vectorStore;

    private final ChatSessionService chatSessionService;

    private static final Map<String, Boolean> GENERATE_STATUS = new ConcurrentHashMap<>();

    private static final ChatEventVO STOP_EVENT = ChatEventVO.builder().eventType(ChatEventTypeEnum.STOP.getValue()).build();

    /**
     * 聊天
     *
     * @param question  问题
     * @param sessionId 会话id
     * @return 回答内容
     */
    @Override
    public Flux<ChatEventVO> chat(String question, String sessionId) {
        Long userId = UserContext.getUser();
        String requestId = IdUtil.fastSimpleUUID();
        StringBuilder stringBuilder = new StringBuilder();
        chatSessionService.update(sessionId,question, userId);
        return dashScopeChatClient.prompt()
                .system(promptSystemSpec -> promptSystemSpec
                        .text(systemPromptConfig.getChatSystemMessage().get())
                        .param("now", DateUtil.now()))
                .advisors(advisorSpec -> advisorSpec
                        .advisors(new QuestionAnswerAdvisor(vectorStore, SearchRequest.builder().query("").topK(999).build()))
                        .param(AbstractChatMemoryAdvisor.CHAT_MEMORY_CONVERSATION_ID_KEY, sessionId))
                .user(question)
                .toolContext(Map.of(Constant.REQUEST_ID, requestId,Constant.USER_ID, userId))
                .stream()
                .chatResponse()
                .doFirst(() -> GENERATE_STATUS.put(sessionId, true))
                .doOnComplete(() -> GENERATE_STATUS.remove(sessionId))
                .doOnError(throwable -> GENERATE_STATUS.remove(sessionId))
                .doOnCancel(() -> saveStopHistoryRecord(sessionId, stringBuilder.toString()))
                .takeWhile(s -> GENERATE_STATUS.getOrDefault(sessionId, false))
                .map(response -> {
                    String finishReason = response.getResult().getMetadata().getFinishReason();
                    if (Constant.STOP.equals(finishReason)){
                        String id = response.getMetadata().getId();
                        ToolResultHolder.put(id,Constant.REQUEST_ID,requestId);
                    }
                    String context = response.getResult().getOutput().getText();
                    stringBuilder.append(context);
                    return ChatEventVO.builder()
                            .eventData(context)
                            .eventType(ChatEventTypeEnum.DATA.getValue())
                            .build();
                })
                .concatWith(Flux.defer(() -> {
                    Map<String, Object> map = ToolResultHolder.get(requestId);
                    if (MapUtil.isNotEmpty(map)) {
                        ToolResultHolder.remove(requestId);
                        ChatEventVO paramEvent = ChatEventVO.builder().eventData(map).eventType(ChatEventTypeEnum.PARAM.getValue()).build();
                        return Flux.just(paramEvent, STOP_EVENT);
                    }
                    return Flux.just(STOP_EVENT);
                }));
    }

    /**
     * 停止生成
     *
     * @param sessionId 会话id
     */
    @Override
    public void stop(String sessionId) {
        GENERATE_STATUS.remove(sessionId);
    }

    /**
     * 保存中止的对话
     * @param conversationId
     * @param context
     */
    private void saveStopHistoryRecord(String conversationId, String context) {
        chatMemory.add(conversationId, new AssistantMessage(context));
    }

}
